package com.yj.limit;

import lombok.extern.slf4j.Slf4j;
import org.aspectj.lang.ProceedingJoinPoint;
import org.aspectj.lang.annotation.Around;
import org.aspectj.lang.annotation.Aspect;
import org.aspectj.lang.annotation.Pointcut;
import org.aspectj.lang.reflect.MethodSignature;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.data.redis.core.RedisTemplate;
import org.springframework.stereotype.Component;
import org.springframework.util.StringUtils;
import org.springframework.web.context.request.RequestContextHolder;
import org.springframework.web.context.request.ServletRequestAttributes;

import javax.servlet.http.HttpServletRequest;
import java.lang.reflect.Method;
import java.util.concurrent.TimeUnit;

@Component
@Aspect
@Slf4j
public class RequestLimitAop {

    @Autowired
    private RedisTemplate<String,String> redisTemplate;

    @Pointcut("@annotation(com.yj.limit.RequestLimit)")
    public void pointcut() {}

    @Around("pointcut()")
    public Object around(ProceedingJoinPoint joinPoint) throws Throwable {
        MethodSignature methodSignature = (MethodSignature) joinPoint.getSignature();
        Method method = methodSignature.getMethod();
        // 获取当方法上的注解
        RequestLimit requestLimit = method.getAnnotation(RequestLimit.class);

        String key = requestLimit.key();
        int limit = requestLimit.limit();
        int timeout = requestLimit.period();

        // 组装rateLimitKey
        //String rateLimitKey = key + getClientAddress();
        String rateLimitKey = key;

        // 获取当前时间
        long now = System.currentTimeMillis();
        // 删除超时的key
        long expireTime = now - timeout;
        // 清理过期的记录
        redisTemplate.opsForZSet().removeRangeByScore(rateLimitKey, 0, expireTime);

        long requestTime = redisTemplate.opsForZSet().count(rateLimitKey,Double.MIN_VALUE, Double.MAX_VALUE);
        // 判断请求是否超过限流阈值
        if (requestTime <= limit){
            redisTemplate.opsForZSet().add(rateLimitKey, String.valueOf(now), now);
            redisTemplate.expire(rateLimitKey, timeout, TimeUnit.SECONDS);
            return joinPoint.proceed();
        }else{
            log.warn("请求频率超过限制，请稍后再试！");
            throw new RequestLimitException("请求频率超过限制，请稍后再试！");
        }
    }


    public String getClientAddress(){

        // 获取request对象
        ServletRequestAttributes attributes = (ServletRequestAttributes) RequestContextHolder.getRequestAttributes();
        HttpServletRequest request = attributes.getRequest();
        // 获取客户端IP地址
        String ip = request.getHeader("x-forwarded-for");
        if (ip == null || ip.length() == 0 || "unknown".equalsIgnoreCase(ip)) {
            ip = request.getHeader("Proxy-Client-IP");
        }
        if (ip == null || ip.length() == 0 || "unknown".equalsIgnoreCase(ip)) {
            ip = request.getHeader("WL-Proxy-Client-IP");
        }
        if (ip == null || ip.length() == 0 || "unknown".equalsIgnoreCase(ip)) {
            ip = request.getRemoteAddr();
        }
        return ip;
    }
}
